import numpy as np
import h5py
from scipy.optimize import minimize
from tools.Cal_Z import Cal_Z
from tools.read_hdf5_find_powers import read_hdf5_find_powers
from tools.mag_phase_to_IQ import mag_phase_to_IQ
from tools.calculate_deri import calculate_deri
from tools.Fdecay_fhalf import Fdecay_fhalf
from concurrent.futures import ProcessPoolExecutor
import multiprocessing
import time
import os
import glob
import argparse

# 系统参数定义
RQ = 8893      # 电阻-品质因数乘积，单位：Ω
Ql = 76.7      # 外部品质因数，无量纲
beta = 1.2     # 耦合系数，无量纲
Fs = 117.402e6 / 14  # 采样率，单位：Hz

def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='计算腔体校准参数a')
    parser.add_argument('--all', action='store_true', help='处理所有时间戳')
    parser.add_argument('--start', type=int, help='起始索引（包含）')
    parser.add_argument('--end', type=int, help='结束索引（包含）')
    return parser.parse_args()

def get_timestamp_count(filename):
    """获取HDF5文件中可用的时间戳数量"""
    with h5py.File(filename, 'r') as f:
        # 使用第一个非axis键来检查时间戳数量
        for key in f.keys():
            if 'XAxis' not in key:
                return len(f[key])
    return 0

def calculate_y_a(a, Z, Vc, Vf, Vr):
    w_half = Fdecay_fhalf(Vc, Vf, Fs)
    a_complex = complex(a[0], a[1])
    Vf_cal = a_complex * Vf + a_complex / Z * Vr
    rho = 2 * np.abs((beta / (1 + beta)) * Vf_cal)
    return w_half * (rho * np.cos(np.angle(Vc) - np.angle(Vf_cal)) - np.abs(Vc))

def error_function(a, Z, Vc, Vf, Vr):
    dVc_dt = calculate_deri(np.abs(Vc), Fs)
    y_a = calculate_y_a(a, Z, Vc, Vf, Vr)
    return np.sum((dVc_dt - y_a) ** 2)

def save_data(filename, results):
    """保存数据到HDF5文件"""
    # 确保输出目录存在
    os.makedirs(os.path.dirname(filename), exist_ok=True)
    
    with h5py.File(filename, 'w') as f:
        for idx, result in enumerate(results):
            if result is not None:
                a_opt, Vc, Vf, Vr = result
                
                # 创建索引数据集
                index_group = f.create_group(f'index_{idx+1}')
                
                # 保存a参数
                index_group.create_dataset('a_real', data=a_opt.real)
                index_group.create_dataset('a_imag', data=a_opt.imag)
                
                # 保存电压数据
                index_group.create_dataset('Vc_real', data=Vc.real)
                index_group.create_dataset('Vc_imag', data=Vc.imag)
                index_group.create_dataset('Vf_real', data=Vf.real)
                index_group.create_dataset('Vf_imag', data=Vf.imag)
                index_group.create_dataset('Vr_real', data=Vr.real)
                index_group.create_dataset('Vr_imag', data=Vr.imag)

def process_single_index(args):
    """处理单个索引的数据"""
    input_filename, file_index = args
    try:
        # 读取数据
        raw_powers = read_hdf5_find_powers(input_filename, file_index)
        if not raw_powers:  # 如果返回None或空字典
            return None
            
        voltages_IQ = mag_phase_to_IQ(raw_powers, RQ, Ql, beta)

        # 转换为复数形式
        Vc = voltages_IQ['Vcav_I'] + 1j * voltages_IQ['Vcav_Q']
        Vf = voltages_IQ['Vfor_I'] + 1j * voltages_IQ['Vfor_Q']
        Vr = voltages_IQ['Vrefl_I'] + 1j * voltages_IQ['Vrefl_Q']

        # 计算参数
        Z = Cal_Z(Vf, Vr)
        result = minimize(error_function, [0, 0.0], args=(Z, Vc, Vf, Vr), method='Nelder-Mead')
        a_opt = complex(result.x[0], result.x[1])

        print(f"索引 {file_index}:\ta = {a_opt:.6f}")
        
        return a_opt, Vc, Vf, Vr
        
    except Exception as e:
        print(f"处理索引 {file_index} 时出错: {str(e)}")
        return None

def process_file(input_filename, start_idx=None, end_idx=None, process_all=False):
    """使用多进程处理文件的所有索引数据"""
    # 获取可用的时间戳数量
    timestamp_count = get_timestamp_count(input_filename)
    if timestamp_count == 0:
        print("错误：无法获取时间戳数量")
        return []
    
    print(f"文件中包含 {timestamp_count} 个时间戳")
    
    # 确定要处理的索引范围
    if process_all:
        start_idx = 1
        end_idx = timestamp_count
    else:
        # 如果没有指定范围，使用默认值（1-10）
        start_idx = start_idx if start_idx is not None else 1
        end_idx = end_idx if end_idx is not None else min(10, timestamp_count)
        
        # 检查索引范围是否有效
        if start_idx < 1:
            print(f"警告：起始索引 {start_idx} 无效，使用 1")
            start_idx = 1
        if end_idx > timestamp_count:
            print(f"警告：结束索引 {end_idx} 超出范围，使用 {timestamp_count}")
            end_idx = timestamp_count
        if start_idx > end_idx:
            print(f"错误：起始索引 {start_idx} 大于结束索引 {end_idx}")
            return []
    
    print(f"处理索引范围: {start_idx} - {end_idx}")
    
    # 准备进程池
    num_cores = multiprocessing.cpu_count()
    print(f"使用 {num_cores} 个CPU核心进行并行计算")
    
    # 准备参数
    args_list = [(input_filename, i) for i in range(start_idx, end_idx + 1)]
    
    # 使用进程池执行计算
    with ProcessPoolExecutor(max_workers=num_cores) as executor:
        results = list(executor.map(process_single_index, args_list))
    
    return results

def main():
    # 在这里设置处理参数
    # 方式1：处理所有时间戳
    process_all = True
    
    # 方式2：处理指定范围的时间戳（如果process_all为False且start_idx和end_idx都有值，则使用这个范围）
    start_idx = None  # 例如：设置为5表示从第5个开始
    end_idx = None    # 例如：设置为15表示到第15个结束
    
    # 如果使用命令行参数，则解析命令行参数（可以注释掉这部分来完全使用上面的设置）
    args = parse_args()
    if args.all or args.start is not None or args.end is not None:
        process_all = args.all
        start_idx = args.start
        end_idx = args.end
    
    # 获取输入目录中的所有HDF5文件
    input_dir = r'E:\MLdata\AFFon'
    hdf5_files = glob.glob(os.path.join(input_dir, '*.hdf5'))
    if not hdf5_files:
        print("错误：在指定目录中未找到HDF5文件")
        return
    
    print(f"找到 {len(hdf5_files)} 个HDF5文件")
    print("-" * 80)
    
    # 处理每个文件
    for input_filename in hdf5_files:
        print(f"\n处理文件: {input_filename}")
        
        # 设置输出文件路径，使用输入文件的名字
        output_dir = r'E:\MLdata\AFFon\cali_data'
        input_basename = os.path.basename(input_filename)  # 获取输入文件名（带扩展名）
        output_basename = os.path.splitext(input_basename)[0] + '_cali.h5'  # 修改扩展名
        output_filename = os.path.join(output_dir, output_basename)
        
        print("开始并行处理数据...")
        print("-" * 40)
        
        # 记录开始时间
        start_time = time.time()
        
        # 使用多进程处理文件
        results = process_file(input_filename, 
                             start_idx=start_idx,
                             end_idx=end_idx,
                             process_all=process_all)
        
        if not results:
            print("错误：未能获取任何有效结果")
            continue
        
        # 保存所有数据
        save_data(output_filename, results)
        
        # 计算总用时
        total_time = time.time() - start_time
        
        print("-" * 40)
        print(f"数据已保存到文件: {output_filename}")
        print(f"处理时间: {total_time:.2f} 秒")
    
    print("\n所有文件处理完成！")

if __name__ == '__main__':
    main() 